import torch
import torch.nn as nn


class CrossEntropyLossPlugin:
    """与MiniMind保持一致的交叉熵损失插件"""
    
    def __init__(self):
        self.loss_fct = nn.CrossEntropyLoss(reduction='none')
    
    def compute_loss(self, logits, labels, loss_mask=None):
        """
        计算损失，与MiniMind保持一致
        """
        # 确保logits和labels形状匹配
        loss = self.loss_fct(
            logits.view(-1, logits.size(-1)),
            labels.view(-1)
        ).view(labels.size())
        
        # 如果提供了loss_mask，则应用它
        if loss_mask is not None:
            loss = (loss * loss_mask).sum() / loss_mask.sum()
        else:
            # 如果没有提供loss_mask，则计算整个序列的平均损失
            loss = loss.mean()
            
        return loss